package cn.edu.szu;

import java.io.File;

public class AESUtils {

    // 128位AES加密轮数为10
    private final int NR = 10;

    // 128位AES密钥长度
    private final int NK = 4;

    // 128位AES分组长度为4
    private final int NB = 4;

    // 密钥
    private final byte[] key;

    // 轮密钥
    private final byte []roundKey;

    // S盒
    private final byte[] sbox = {
            /* 0     1     2     3     4     5     6     7     8     9     A     B     C     D     E     F  */
            0x63, 0x7c, 0x77, 0x7b, (byte) 0xf2, 0x6b, 0x6f, (byte)0xc5, 0x30, 0x01, 0x67, 0x2b, (byte)0xfe, (byte)0xd7, (byte)0xab, (byte)0x76,
            (byte) 0xca, (byte) 0x82, (byte) 0xc9, 0x7d, (byte) 0xfa, 0x59, 0x47, (byte) 0xf0, (byte) 0xad, (byte) 0xd4, (byte) 0xa2, (byte) 0xaf, (byte) 0x9c, (byte)0xa4, 0x72, (byte)0xc0,
            (byte) 0xb7, (byte) 0xfd, (byte) 0x93, 0x26, 0x36, 0x3f, (byte) 0xf7, (byte) 0xcc, 0x34, (byte) 0xa5, (byte) 0xe5, (byte) 0xf1, 0x71, (byte) 0xd8, 0x31, 0x15,
            0x04, (byte) 0xc7, 0x23, (byte) 0xc3, 0x18, (byte) 0x96, 0x05, (byte) 0x9a, 0x07, 0x12, (byte) 0x80, (byte) 0xe2, (byte) 0xeb, 0x27, (byte) 0xb2, 0x75,
            0x09, (byte) 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, (byte) 0xa0, 0x52, 0x3b, (byte) 0xd6, (byte) 0xb3, 0x29, (byte) 0xe3, 0x2f, (byte) 0x84,
            0x53, (byte) 0xd1, 0x00, (byte) 0xed, 0x20, (byte) 0xfc, (byte) 0xb1, 0x5b, 0x6a, (byte) 0xcb, (byte) 0xbe, 0x39, 0x4a, 0x4c, 0x58, (byte) 0xcf,
            (byte) 0xd0, (byte) 0xef, (byte) 0xaa, (byte) 0xfb, 0x43, 0x4d, 0x33, (byte) 0x85, 0x45, (byte) 0xf9, 0x02, 0x7f, 0x50, 0x3c, (byte) 0x9f, (byte) 0xa8,
            0x51, (byte) 0xa3, 0x40, (byte) 0x8f, (byte) 0x92, (byte) 0x9d, 0x38, (byte) 0xf5, (byte) 0xbc, (byte) 0xb6, (byte) 0xda, 0x21, 0x10, (byte) 0xff, (byte) 0xf3, (byte) 0xd2,
            (byte) 0xcd, 0x0c, 0x13, (byte) 0xec, 0x5f, (byte) 0x97, 0x44, 0x17, (byte) 0xc4, (byte) 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
            0x60, (byte) 0x81, 0x4f, (byte) 0xdc, 0x22, 0x2a, (byte) 0x90, (byte) 0x88, 0x46, (byte) 0xee, (byte) 0xb8, 0x14, (byte) 0xde, 0x5e, 0x0b, (byte) 0xdb,
            (byte) 0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, (byte) 0xc2, (byte) 0xd3, (byte) 0xac, 0x62, (byte) 0x91, (byte) 0x95, (byte) 0xe4, 0x79,
            (byte) 0xe7, (byte) 0xc8, 0x37, 0x6d, (byte) 0x8d, (byte) 0xd5, 0x4e, (byte) 0xa9, 0x6c, 0x56, (byte) 0xf4, (byte) 0xea, 0x65, 0x7a, (byte) 0xae, 0x08,
            (byte) 0xba, 0x78, 0x25, 0x2e, 0x1c, (byte) 0xa6, (byte) 0xb4, (byte) 0xc6, (byte) 0xe8, (byte) 0xdd, 0x74, 0x1f, 0x4b, (byte) 0xbd, (byte) 0x8b, (byte) 0x8a,
            0x70, 0x3e, (byte) 0xb5, 0x66, 0x48, 0x03, (byte) 0xf6, 0x0e, 0x61, 0x35, 0x57, (byte) 0xb9, (byte) 0x86, (byte) 0xc1, 0x1d, (byte) 0x9e,
            (byte) 0xe1, (byte) 0xf8, (byte) 0x98, 0x11, 0x69, (byte) 0xd9, (byte) 0x8e, (byte) 0x94, (byte) 0x9b, 0x1e, (byte) 0x87, (byte) 0xe9, (byte) 0xce, 0x55, 0x28, (byte) 0xdf,
            (byte) 0x8c, (byte) 0xa1, (byte) 0x89, 0x0d, (byte) 0xbf, (byte) 0xe6, 0x42, 0x68, 0x41, (byte) 0x99, 0x2d, 0x0f, (byte) 0xb0, 0x54, (byte) 0xbb, 0x16
    };

    // 逆S盒
    private final byte[] inv_sbox = {
            /* 0     1     2     3     4     5     6     7     8     9     A     B     C     D     E     F  */
            0x52, 0x09, 0x6a, (byte) 0xd5, 0x30, 0x36, (byte) 0xa5, 0x38, (byte) 0xbf, 0x40, (byte) 0xa3, (byte) 0x9e, (byte) 0x81, (byte) 0xf3, (byte) 0xd7, (byte) 0xfb,
            0x7c, (byte) 0xe3, 0x39, (byte) 0x82, (byte) 0x9b, 0x2f, (byte) 0xff, (byte) 0x87, 0x34, (byte) 0x8e, 0x43, 0x44, (byte) 0xc4, (byte) 0xde, (byte) 0xe9, (byte) 0xcb,
            0x54, 0x7b, (byte) 0x94, 0x32, (byte) 0xa6, (byte) 0xc2, 0x23, 0x3d, (byte) 0xee, 0x4c, (byte) 0x95, 0x0b, 0x42, (byte) 0xfa, (byte) 0xc3, 0x4e,
            0x08, 0x2e, (byte) 0xa1, 0x66, 0x28, (byte) 0xd9, 0x24, (byte) 0xb2, 0x76, 0x5b, (byte) 0xa2, 0x49, 0x6d, (byte) 0x8b, (byte) 0xd1, 0x25,
            0x72, (byte) 0xf8, (byte) 0xf6, 0x64, (byte) 0x86, 0x68, (byte) 0x98, 0x16, (byte) 0xd4, (byte) 0xa4, 0x5c, (byte) 0xcc, 0x5d, 0x65, (byte) 0xb6, (byte) 0x92,
            0x6c, 0x70, 0x48, 0x50, (byte) 0xfd, (byte) 0xed, (byte) 0xb9, (byte) 0xda, 0x5e, 0x15, 0x46, 0x57, (byte) 0xa7, (byte) 0x8d, (byte) 0x9d, (byte) 0x84,
            (byte) 0x90, (byte)0xd8, (byte) 0xab, 0x00, (byte) 0x8c, (byte) 0xbc, (byte) 0xd3, 0x0a, (byte) 0xf7, (byte) 0xe4, 0x58, 0x05, (byte) 0xb8, (byte) 0xb3, 0x45, 0x06,
            (byte) 0xd0, 0x2c, 0x1e, (byte) 0x8f, (byte) 0xca, 0x3f, 0x0f, 0x02, (byte) 0xc1, (byte) 0xaf, (byte) 0xbd, 0x03, 0x01, 0x13, (byte) 0x8a, 0x6b,
            0x3a, (byte) 0x91, 0x11, 0x41, 0x4f, 0x67, (byte) 0xdc, (byte) 0xea, (byte) 0x97, (byte) 0xf2, (byte) 0xcf, (byte) 0xce, (byte) 0xf0, (byte) 0xb4, (byte) 0xe6, 0x73,
            (byte) 0x96, (byte) 0xac, 0x74, 0x22, (byte) 0xe7, (byte) 0xad, 0x35, (byte) 0x85, (byte) 0xe2, (byte) 0xf9, 0x37, (byte) 0xe8, 0x1c, 0x75, (byte) 0xdf, 0x6e,
            0x47, (byte) 0xf1, 0x1a, 0x71, 0x1d, 0x29, (byte) 0xc5, (byte) 0x89, 0x6f, (byte) 0xb7, 0x62, 0x0e, (byte) 0xaa, 0x18, (byte) 0xbe, 0x1b,
            (byte) 0xfc, 0x56, 0x3e, 0x4b, (byte) 0xc6, (byte) 0xd2, 0x79, 0x20, (byte) 0x9a, (byte) 0xdb, (byte) 0xc0, (byte) 0xfe, 0x78, (byte) 0xcd, 0x5a, (byte) 0xf4,
            0x1f, (byte) 0xdd, (byte) 0xa8, 0x33, (byte) 0x88, 0x07, (byte) 0xc7, 0x31, (byte) 0xb1, 0x12, 0x10, 0x59, 0x27, (byte) 0x80, (byte) 0xec, 0x5f,
            0x60, 0x51, 0x7f, (byte) 0xa9, 0x19, (byte) 0xb5, 0x4a, 0x0d, 0x2d, (byte) 0xe5, 0x7a, (byte) 0x9f, (byte) 0x93, (byte) 0xc9, (byte) 0x9c, (byte) 0xef,
            (byte) 0xa0, (byte) 0xe0, 0x3b, 0x4d, (byte) 0xae, 0x2a, (byte) 0xf5, (byte) 0xb0, (byte) 0xc8, (byte) 0xeb, (byte) 0xbb, 0x3c, (byte) 0x83, 0x53, (byte) 0x99, 0x61,
            0x17, 0x2b, 0x04, 0x7e, (byte) 0xba, 0x77, (byte) 0xd6, 0x26, (byte) 0xe1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0c, 0x7d
    };

    private final byte []Rcon = {
            (byte) 0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, (byte) 0x80, 0x1b, 0x36
    };

    private final byte []y = {2, 3, 1, 1,  1, 2, 3, 1,
            1, 1, 2, 3,  3, 1, 1, 2};

    private final byte []inv_y = {0x0e, 0x0b, 0x0d, 0x09,  0x09, 0x0e, 0x0b, 0x0d,
            0x0d, 0x09, 0x0e, 0x0b,  0x0b, 0x0d, 0x09, 0x0e};

    // 构造方法保证外部无法实例化
    public AESUtils(String key_string) {
        // 初始化密钥
        key = hexStringToByteArray(key_string);

        // 初始化轮密钥
        roundKey = new byte[(NR + 1) * 16];
        keyExpansion(roundKey);

    }

    // 初始化文件和密钥值
    public void init() {

    }

    // 密钥拓展算法
    private void keyExpansion(byte []roundKey) {
        byte []temp = new byte[4];

        // 第0轮的密钥是key自身
        for (int i = 0; i < NK; ++i) {
            roundKey[i * 4] = key[i * 4];
            roundKey[i * 4 + 1] = key[i * 4 + 1];
            roundKey[i * 4 + 2] = key[i * 4 + 2];
            roundKey[i * 4 + 3] = key[i * 4 + 3];
        }

        // 新一轮密钥依赖前一轮的密钥
        for (int i = NK; i < NB * (NR + 1); ++i) {
            {
                int k = (i - 1) * 4;
                temp[0] = roundKey[k];
                temp[1] = roundKey[k + 1];
                temp[2] = roundKey[k + 2];
                temp[3] = roundKey[k + 3];
            }
            if (i % NK == 0) {
                // RotWord()函数
                {
                    byte tmp = temp[0];
                    temp[0] = temp[1];
                    temp[1] = temp[2];
                    temp[2] = temp[3];
                    temp[3] = tmp;
                }
                // Subword()函数
                {
                    temp[0] = sbox[temp[0] & 0xff];
                    temp[1] = sbox[temp[1] & 0xff];
                    temp[2] = sbox[temp[2] & 0xff];
                    temp[3] = sbox[temp[3] & 0xff];
                }

                temp[0] = (byte) (temp[0] ^ Rcon[i/NK]);
            }
            int j = i * 4;
            int k = (i - NK) * 4;
            roundKey[j] = (byte) (roundKey[k] ^ temp[0]);
            roundKey[j + 1] = (byte) (roundKey[k + 1] ^ temp[1]);
            roundKey[j + 2] = (byte) (roundKey[k + 2] ^ temp[2]);
            roundKey[j + 3] = (byte) (roundKey[k + 3] ^ temp[3]);
        }
    }

    // 加密
    public void encrypt(byte []state) {

        int round = 0;
        add_round_key(round, state);

        for (round=1; round<NR; ++round) {
            sub_bytes(state);
            shift_rows(state);
            mix_columns(state);
            add_round_key(round, state);
        }

        sub_bytes(state);
        shift_rows(state);
        add_round_key(NR, state);
    }

    // 解密
    public void decrypt(byte []state) {

        int round = NR;

        add_round_key(round, state);
        inv_sub_bytes(state);
        inv_shift_rows(state);

        for(round = (NR - 1); round > 0; --round) {
            add_round_key(round, state);
            inv_mix_columns(state);
            inv_sub_bytes(state);
            inv_shift_rows(state);
        }

        add_round_key(round, state);
    }

    // 字节代换
    private void sub_bytes(byte []state) {
        for (int i=0; i<NB; ++i) {
            for (int j=0; j<NB; ++j) {
                state[i * 4 + j] = sbox[state[i * 4 + j] & 0xff];
            }
        }
    }

    // 逆字节代换
    private void inv_sub_bytes(byte []state) {
        for (int i=0; i<NB; ++i) {
            for (int j=0; j<NB; ++j) {
                state[i * 4 + j] = inv_sbox[state[i * 4 + j] & 0xff];
            }
        }
    }

    // 行移位，循环展开以提高效率
    private void shift_rows(byte []state) {
        byte tmp;

        // 第1行
        tmp = state[1];
        state[1] = state[5];
        state[5] = state[9];
        state[9] = state[13];
        state[13] = tmp;

        // 第2行
        tmp = state[2];        state[2] = state[10];        state[10] = tmp;
        tmp = state[6];        state[6] = state[14];        state[14] = tmp;

        // 第3行
        tmp = state[3];
        state[3] = state[15];
        state[15] = state[11];
        state[11] = state[7];
        state[7] = tmp;
    }

    // 逆向行移位
    private void inv_shift_rows(byte []state) {
        byte tmp;

        // 第1行
        tmp = state[13];
        state[13] = state[9];
        state[9] = state[5];
        state[5] = state[1];
        state[1] = tmp;

        // 第2行
        tmp = state[14];        state[14] = state[6];        state[6] = tmp;
        tmp = state[10];        state[10] = state[2];        state[2] = tmp;

        // 第3行
        tmp = state[7];
        state[7] = state[11];
        state[11] = state[15];
        state[15] = state[3];
        state[3] = tmp;
    }

    // 有限域上的乘法 GF(2^8)
    private byte mul(byte a, byte b) {
        byte p = 0;

        for (int counter = 0; counter < 8; counter++) {
            if ((b & 1) != 0) {
                p ^= a;
            }

            boolean hi_bit_set = (a & 0x80) != 0;
            a <<= 1;
            if (hi_bit_set) {
                a ^= 0x1B; /* x^8 + x^4 + x^3 + x + 1 */
            }
            b >>= 1;
        }

        return p;
    }

    // 列混合
    private void mix_columns(byte []state) {
        byte []s = new byte[4];
        for (int i=0; i<NB; ++i) {
            for (int r=0; r<4; ++r) {
                s[r] = 0;
                // 矩阵乘法
                for (int j=0; j<4; ++j) {
                    // 减少一点函数调用，提速
                    if (y[r * 4 + j] == 1) {
                        s[r] = (byte) (s[r] ^ state[i * 4 + j]);
                        continue;
                    }
                    s[r] = (byte) (s[r] ^ mul(state[i * 4 + j], y[r * 4 + j]));
                }
            }
            // 赋值
            System.arraycopy(s, 0, state, i * 4, 4);
        }
    }

    // 逆向列混合
    private void inv_mix_columns(byte []state) {
        byte []s = new byte[4];

        for (int i=0; i<NB; ++i) {
            for (int r=0; r<4; ++r) {
                s[r] = 0;
                // 矩阵乘法
                for (int j=0; j<4; ++j) {
                    s[r] = (byte) (s[r] ^ mul(state[i * 4 + j], inv_y[r * 4 + j]));
                }
            }
            // 赋值
            System.arraycopy(s, 0, state, i * 4, 4);
        }
    }

    // 轮密钥加
    private void add_round_key(int round, byte []state) {
        for (int i=0; i<NB * 4; ++i) {
            state[i] = (byte) (state[i] ^ roundKey[round * NB * 4 + i]);
        }
    }

    // 将字符串转化为byte[]
    private byte[] hexStringToByteArray(String s) {
        int length = s.length();
        byte[] state = new byte[length / 2];
        for (int i=0; i<length; i+=2) {
            state[i/2] = (byte)((Character.digit(s.charAt(i), 16) << 4) + Character.digit(s.charAt(i+1), 16) );
        }
        return state;
    }

    private void print(byte[] data) {
        for (int i=0; i<NB * 4; ++i) {
            System.out.print(Integer.toHexString((data[i]) & 0xff) + " ");
        }
        System.out.println();
    }

    private void print_round(int round, byte[] data) {
        System.out.print("roundKey:");
        for (int i=round*16; i<round*16+16; ++i) {
            System.out.print(Integer.toHexString((data[i]) & 0xff) + " ");
        }
        System.out.println();
    }
}